package com.kongjs.common.datasource.component;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.jdbc.datasource.lookup.DataSourceLookup;
import org.springframework.jdbc.datasource.lookup.JndiDataSourceLookup;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;


@SuppressWarnings("all")
public class DynamicDataSource extends AbstractRoutingDataSource {
    public static final String DATA_SOURCE_PREFIX = "dynamicDataSource_";
    private final DataSourceKeyBean dataSourceKeyBean;
    private String dataSourcePrefix;
    private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();

    public DynamicDataSource(DataSourceKeyBean dataSourceKeyBean) {
        this.dataSourceKeyBean = dataSourceKeyBean;
    }

    @Override
    protected Object determineCurrentLookupKey() {
        Object lookupKey = dataSourceKeyBean.determineCurrentLookupKey();
        return this.resolveSpecifiedLookupKey(lookupKey);
    }

    @Override
    protected Object resolveSpecifiedLookupKey(Object lookupKey) {
        if (lookupKey instanceof String) {
            if (lookupKey.toString().startsWith(this.getDataSourcePrefix())) {
                return lookupKey;
            } else {
                return this.getDataSourcePrefix() + lookupKey;
            }
        }
        return lookupKey;
    }

    @Override
    protected DataSource resolveSpecifiedDataSource(Object dataSourceObject) throws IllegalArgumentException {
        if (dataSourceObject instanceof DataSource dataSource) {
            return dataSource;
        } else if (dataSourceObject instanceof String dataSourceName && dataSourceName.startsWith(this.getDataSourcePrefix())) {
            return this.getDataSourceLookup().getDataSource(dataSourceName);
        } else {
            throw new IllegalArgumentException("Illegal data source value - only [javax.sql.DataSource] and String supported: " + String.valueOf(dataSourceObject));
        }
    }

    public void add(Object lookupKey, DataSource ds) {
        Object key = this.resolveSpecifiedLookupKey(lookupKey);
        Map<Object, Object> targetDataSources = new HashMap<>(this.getResolvedDataSources());
        targetDataSources.put(key, ds);
        this.setDataSource(targetDataSources);
    }

    public void remove(Object lookupKey) {
        Object key = this.resolveSpecifiedLookupKey(lookupKey);
        Map<Object, Object> targetDataSources = new HashMap<>(this.getResolvedDataSources());
        targetDataSources.remove(key);
        this.setDataSource(targetDataSources);
    }

    public void setDataSource(Map<Object, Object> targetDataSources) {
        this.setTargetDataSources(targetDataSources);
        this.afterPropertiesSet();
    }

    public String getDataSourcePrefix() {
        if (StringUtils.hasText(dataSourcePrefix)) {
            return dataSourcePrefix;
        }
        return DATA_SOURCE_PREFIX;
    }

    public void setDataSourcePrefix(String dataSourcePrefix) {
        this.dataSourcePrefix = dataSourcePrefix;
    }

    public DataSourceLookup getDataSourceLookup() {
        return dataSourceLookup;
    }

    @Override
    public void setDataSourceLookup(DataSourceLookup dataSourceLookup) {
        this.dataSourceLookup = dataSourceLookup;
    }
}
